from langchain.chains import create_sql_query_chain
from langchain.prompts import PromptTemplate
from boat_nl2sql.utiles import extract_sql_from_markdown, load_table_info_from_json
from boat_nl2sql.oracle_connection import get_db_connection, FewShotPromptTemplate
from boat_nl2sql.generate_echarts import visualize_sql_result, analyze_data_for_visualization, convert_sql_result_to_dataframe

# 在这里导入您的 LLM
from boat_nl2sql.llm_module import SiliconFlow
llm = SiliconFlow()

def generate_analysis(question):

    print("\n===== 测试 SQL 查询链 =====")

    # 加载表信息
    full_table_info = load_table_info_from_json("boat_nl2sql/oracle_table_info.json")
    include_tables = [
        'BAYONET_BUSSINESS.BOAT_WARNING',
        'BAYONET_BUSSINESS.OFF_SITE_CASE',
        'BAYONET_BUSSINESS.ZD_WARNING_TYPE',
        'BAYONET_DYNAMIC.VIDEO_ANALYSIS',
        'BAYONET_DYNAMIC.DATAFUSION',
        'BAYONET_BASICS.SYS_BAYONET'
    ]

    # 获取数据库连接
    db = get_db_connection(
        sample_rows=2,
        custom_table_info=full_table_info,
        include_tables=include_tables
    )

    # 配置提示模板
    SQL_PROMPT = PromptTemplate.from_template(
        """你是一个SQL专家。根据以下表结构：  
        {table_info}  

        问题：{input}  
        请严格按以下规则响应：  
        1. 只输出SQL代码  
        2. 使用标准Oracle SQL语法  
        3. 不要包含分号结尾  
        4. 最多返回{top_k}条记录  
        5. Oracle中使用ROWNUM来限制返回记录数  

        SQL查询："""
    )

    # 配置few_shot提示模板
    examples = [
        {
            "input": "查询名称中包含'船闸'的所有卡口基本信息，按名称排序。",
            "query": "SELECT ID, NAME, UNIT_ID, DES, PHOTOHTTP  FROM BAYONET_BASICS.SYS_BAYONET  WHERE NAME LIKE '%船闸%'  ORDER BY NAME",
        },
        {
            "input": "统计一月份各卡口点的船舶通过数量，按照通过船舶数量降序排列，只显示通过量前10的卡口。",
            "query": "SELECT * FROM (SELECT b.NAME AS 卡口名称, COUNT(d.CODE) AS 通过船舶数量  FROM BAYONET_BASICS.SYS_BAYONET b  JOIN BAYONET_DYNAMIC.DATAFUSION d ON b.ID = d.BAYONET_ID  WHERE d.PASSTIME >= TRUNC(SYSDATE, 'YYYY') AND d.PASSTIME < ADD_MONTHS(TRUNC(SYSDATE, 'YYYY'), 1) GROUP BY b.NAME  ORDER BY 通过船舶数量 DESC) WHERE ROWNUM <= 10",
        },
        {
            "input": "查询系统中各类型预警的发生次数和占比，按发生次数从高到低排序。",
            "query": "WITH warning_counts AS (  SELECT t.NAME AS 预警类型,   COUNT(w.CODE) AS 预警次数  FROM BAYONET_BUSSINESS.ZD_WARNING_TYPE t  JOIN BAYONET_BUSSINESS.BOAT_WARNING w ON t.ID = w.WARNING_TYPE  GROUP BY t.NAME  )  SELECT    预警类型,预警次数, ROUND(预警次数 * 100 / SUM(预警次数) OVER(), 2) || '%' AS 占比 FROM warning_counts ORDER BY 预警次数 DESC "
        },
        {
            "input": "分析一下2025年1月油墩港(东大盈船闸)的流量组成。",
            "query": """SELECT 
        COUNT(da.CODE) AS TOTAL_FLOW,          -- 总流量
        SUM(CASE WHEN da.DIRECTION = 1 THEN 1 ELSE 0 END) AS UPSTREAM_FLOW, -- 上行流量
        SUM(CASE WHEN da.DIRECTION = 2 THEN 1 ELSE 0 END) AS DOWNSTREAM_FLOW -- 下行流量
    FROM 
        BAYONET_DYNAMIC.DATAFUSION da
    LEFT JOIN BAYONET_BASICS.SYS_BAYONET sb
        ON da.BAYONET_ID = sb.ID
    WHERE 
        sb.NAME = '油墩港(东大盈船闸)' 
        AND da.PASSTIME >= TO_DATE('2025-01-01 00:00:00', 'YYYY-MM-DD HH24:MI:SS')
        AND da.PASSTIME <= TO_DATE('2025-01-31 23:59:59', 'YYYY-MM-DD HH24:MI:SS');"""
        },
        {
            "input": "查询油墩港(东大盈船闸)2025年 1月份各类预警的案件总数与占比",
            "query": """WITH filtered_data AS (
      SELECT 
        zwt.NAME AS WARNING_TYPE_NAME
      FROM 
        BAYONET_DYNAMIC.DATAFUSION df
      JOIN 
        BAYONET_BUSSINESS.BOAT_WARNING bw 
        ON df.CODE = bw.EVENTCODE
      JOIN 
        BAYONET_BUSSINESS.ZD_WARNING_TYPE zwt  
        ON bw.WARNING_TYPE = zwt.ID  
      JOIN 
        BAYONET_BASICS.SYS_BAYONET sb  -- 关联卡口基础信息表
        ON df.BAYONET_ID = sb.ID       -- 通过ID关联
      WHERE 
        sb.NAME = '油墩港(东大盈船闸)'  -- 直接使用卡口名称过滤
        AND df.PASSTIME >= DATE '2025-01-01' 
        AND df.PASSTIME < DATE '2025-02-01'
    )

    SELECT  WARNING_TYPE_NAME,  COUNT(*) AS record_count,  ROUND(COUNT(*) * 100.0 / SUM(COUNT(*)) OVER (),2) AS percentage FROM filtered_data GROUP BY WARNING_TYPE_NAME ORDER BY  WARNING_TYPE_NAME;
    """
        }

    ]

    few_shot_example_prompt = PromptTemplate.from_template("输入问题: {input}\n生成SQL: {query}")
    few_shot_prompt = FewShotPromptTemplate(
        examples=examples,
        example_prompt=few_shot_example_prompt,
        prefix="你是一个Oracle SQL专家。请根据以下表结构生成对应问题的Oracle SQL语句，Oracle版本较低(11g及以下)，且尽量符合标准语法,不考虑ISDELETE和ISFINISH字段，如果没有强调年份，默认为2025年。\n\n以下是表结构（table_info）：\n{table_info}\n\n这是一些问题及其对应的SQL样例:",
        suffix="\n---\n问题: {input}\n对应SQL: ",
        input_variables=["input", "table_info"],
    )

    # 创建 SQL 查询链
    generate_query = create_sql_query_chain(llm=llm, db=db, prompt=SQL_PROMPT, k=5)

    few_shot_sql = few_shot_prompt.format(input=question, table_info=db.table_info)

    llm_response = llm.invoke(few_shot_sql)
    print("FewShot 生成的 SQL:", llm_response)

    if "```sql" in llm_response:
        extracted_sql = extract_sql_from_markdown(llm_response)
        print("提取出的 SQL:")
        print(extracted_sql)
    else:
        extracted_sql = llm_response.replace(";", "")

    try:
        # 执行SQL
        sql_result = db.run(extracted_sql)
        print("SQL结果:", sql_result)

        if len(sql_result) > 1:
            # 可视化SQL结果
            output_file = visualize_sql_result(
                sql_result=sql_result,
                query=extracted_sql,
                question=question
            )

            if output_file:
                print(f"可视化已完成，请打开以下文件查看: {output_file}")

                # 生成分析结果
                TEXT_PROMPT = PromptTemplate.from_template(
                    """根据问题{test_question}，SQL查询结果{sql_result}，以及生成的图表类型{chart_type}，  
                    请提供以下内容：  
                    1. 对数据的简明分析  
                    2. 从数据中得出的关键见解  
                    3. 为什么选择了这种图表类型来展示数据  
                    4. 根据数据可能的建议或行动计划  

                    请使用简洁专业的语言回答。"""
                )

                # 从可视化推荐中获取图表类型
                recommendation = analyze_data_for_visualization(
                    convert_sql_result_to_dataframe(sql_result),
                    extracted_sql,
                    question
                )

                chart_type = recommendation["chart_type"]

                # 格式化PromptTemplate为字符串
                formatted_prompt = TEXT_PROMPT.format(
                    test_question=question,
                    sql_result=sql_result,
                    chart_type=chart_type
                )

                analysis_result = llm.invoke(formatted_prompt)
                print("\n数据分析结果:")
                print(analysis_result)
        else:
            TEXT_PROMPT = PromptTemplate.from_template(
                """根据问题{test_question}，SQL查询结果{sql_result}，  
                请提供以下内容：  
                1. 对数据的简明分析  
                2. 从数据中得出的关键见解  
                3. 根据数据可能的建议或行动计划  

                请使用简洁专业的语言回答。"""
            )

            # 格式化PromptTemplate为字符串
            formatted_prompt = TEXT_PROMPT.format(
                test_question=question,
                sql_result=sql_result,
            )

            analysis_result = llm.invoke(formatted_prompt)
            output_file = None
            print("\n数据分析结果:")
            print(analysis_result)

    except Exception as e:
        print(f"执行查询时出错: {e}")

    return output_file, analysis_result


def analyze_sql_data(question, sql_result):
    """基于SQL查询结果生成数据分析报告
    Args:
        question: 需要分析的业务问题
        sql_result: SQL查询结果字符串
        llm: 已初始化的语言模型实例
    Returns:
        包含分析、见解和建议的文本报告
    """
    # 创建分析模板
    analysis_prompt = PromptTemplate.from_template(
        """根据问题{test_question}，SQL查询结果{sql_result}，  
        请提供以下内容：  
        1. 对数据的简明分析  
        2. 从数据中得出的关键见解  
        3. 根据数据可能的建议或行动计划  

        请使用简洁专业的语言回答。"""
    )

    # 格式化PromptTemplate为字符串
    formatted_prompt = analysis_prompt.format(
        test_question=question,
        sql_result=sql_result
    )

    # 执行分析
    analysis_result = llm.invoke(formatted_prompt)

    # 返回分析结果（可选打印）
    print("\n数据分析完成")
    return analysis_result

if __name__ == "__main__":
    question = "统计一月份各卡口点的船舶通过数量，按照通过船舶数量降序排列，只显示通过量前10的卡口。"
    result, outfile = generate_analysis(question)
    print(result)